# Modified from:
#   DiT:  https://github.com/facebookresearch/DiT/blob/main/sample_ddp.py
#   llamagen: https://github.com/FoundationVision/LlamaGen/blob/main/autoregressive/sample/sample_c2i_ddp.py

import torch
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
import torch.nn.functional as F

from tqdm import tqdm
import os
from PIL import Image
import numpy as np
import math
import argparse

import sys 
sys.path.append(os.getcwd())

from autoregressive.models.gpt import GPT_models
from autoregressive.models.generate import generate
from autoregressive.models.mm_utils import direction_index, order_index, label_order_index
from src.IBQ.models.ibqgan import IBQ
from omegaconf import OmegaConf

def load_vqgan_new(config, ckpt_path=None, is_gumbel=False):
	model = IBQ(**config.model.init_args)
	if ckpt_path is not None:
		sd = torch.load(ckpt_path, map_location="cpu")["state_dict"]
		missing, unexpected = model.load_state_dict(sd, strict=True)
	return model.eval()

def custom_to_pil(x):
    x = x.detach().cpu()
    x = torch.clamp(x, -1., 1.)
    x = (x + 1.)/2.
    x *= 255.0
    x = x.permute(0, 2, 3, 1).to(dtype=torch.uint8)
     
    return x

def main(args):
    # Setup PyTorch:
    assert torch.cuda.is_available()
    torch.manual_seed(args.global_seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    torch.set_grad_enabled(False)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # create and load model
    config_file = args.vq_ckpt.replace('imagenet256_16384.ckpt', 'imagenet_ibqgan_16384.yaml')
    configs = OmegaConf.load(config_file)
    dim_z = configs.model.init_args.embed_dim
    vq_model = load_vqgan_new(configs, args.vq_ckpt).to(device)

    # create and load gpt model
    precision = {'none': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}[args.precision]
    latent_size = args.image_size // args.downsample_size
    gpt_model = GPT_models[args.gpt_model](
        vocab_size=args.codebook_size,
        block_size=latent_size ** 2,
        num_classes=args.num_classes,
        cls_token_num=args.cls_token_num,
        vqgan_ckpt=args.vqgan_ckpt,
        rope_base=args.rope_base,
        image_order=args.image_order,
        directions=args.directions,
        condition=args.condition,
        pe_type=args.pe_type,
    ).to(device=device, dtype=precision)

    checkpoint = torch.load(args.gpt_ckpt, map_location="cpu")
    model_weight = checkpoint["model"]

    gpt_model.load_state_dict(model_weight, strict=True)
    gpt_model.eval()
    del checkpoint

    if args.compile:
        print(f"compiling the model...")
        gpt_model = torch.compile(
            gpt_model,
            mode="reduce-overhead",
            fullgraph=True
        ) # requires PyTorch 2.0 (optional)
    else:
        print(f"no model compile") 

    # Create folder to save samples:
    model_string_name = args.gpt_model.replace("/", "-")
    ckpt_string_name = os.path.basename(args.gpt_ckpt).replace(".pth", "").replace(".pt", "")
    folder_name = f"{model_string_name}-{ckpt_string_name}-size-{args.image_size}-" \
                  f"topk-{args.top_k}-topp-{args.top_p}-temperature-{args.temperature}-" \
                  f"cfg-{args.guidance_scale}-{args.guidance_scale_pow}-{args.guidance_start}-seed-{args.global_seed}"
    sample_folder_dir = f"{args.sample_dir}/{folder_name}"

    condition_ids = [torch.randint(0, 999, size=(1,)).item()]

    os.makedirs(sample_folder_dir, exist_ok=True)
    for c in condition_ids:
        os.makedirs(os.path.join(sample_folder_dir, str(c)), exist_ok=True)
    print(f"Saving .png samples at {sample_folder_dir}")

    iterations = 1
    pbar = range(iterations)
    pbar = tqdm(pbar)
    for c in condition_ids:
        total = 0
        for _ in pbar:
            # Sample inputs:
            c_indices = torch.randint(c, c+1, (1,), device=device)
            image_order = label_order_index[args.image_order]

            index_sample = generate(
                gpt_model, c_indices, len(image_order), image_order=order_index[args.image_order],
                label_order=label_order_index[args.image_order],
                guidance_scale=args.guidance_scale,
                guidance_scale_pow=args.guidance_scale_pow,
                guidance_start=args.guidance_start,
                cfg_interval=args.cfg_interval,
                temperature=args.temperature, top_k=args.top_k,
                top_p=args.top_p, sample_logits=True, 
                )
            reorder = [0 for _ in range(latent_size ** 2)]
            for i in range(len(image_order)-1, -1, -1):
                reorder[image_order[i]] = i
            
            index_sample = index_sample[:, reorder]
            bhwc = (len(c_indices), latent_size, latent_size, dim_z)
            samples = vq_model.decode(vq_model.quantize.get_codebook_entry(index_sample, shape=bhwc))
            samples = custom_to_pil(samples).numpy()
            # Save samples to disk as individual .png files
            for i, sample in enumerate(samples):
                index = i + total
                Image.fromarray(sample).save(f"{sample_folder_dir}/{str(c)}/{index:06d}.png")
            total += 1

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--gpt-model", type=str, choices=list(GPT_models.keys()), default="GPT-B")
    parser.add_argument("--gpt-ckpt", type=str, default=None)
    parser.add_argument("--rope-base", type=float, default=10000)
    parser.add_argument("--image-order", type=str, default="default")
    parser.add_argument("--pe-type", type=str, default="default")
    parser.add_argument("--directions", type=int, default=0)
    parser.add_argument("--condition", action='store_true')
    parser.add_argument("--vqgan-ckpt", type=str, default='identity')
    parser.add_argument("--cls-token-num", type=int, default=1, help="max token number of condition input")
    parser.add_argument("--precision", type=str, default='bf16', choices=["none", "fp16", "bf16"]) 
    parser.add_argument("--compile", action='store_true', default=False)
    parser.add_argument("--vq-ckpt", type=str, default=None, help="ckpt path for vq model")
    parser.add_argument("--codebook-size", type=int, default=16384, help="codebook size for vector quantization")
    parser.add_argument("--image-size", type=int, default=256)
    parser.add_argument("--downsample-size", type=int, default=16)
    parser.add_argument("--num-classes", type=int, default=1000)
    parser.add_argument("--cfg-interval", type=float, default=-1)
    parser.add_argument("--guidance-start", type=float, default=1.0)
    parser.add_argument("--guidance-scale", type=float, default=4.7)
    parser.add_argument("--guidance-scale-pow", type=float, default=0.88)
    parser.add_argument("--sample-dir", type=str, default="samples")
    parser.add_argument("--global-seed", type=int, default=0)
    parser.add_argument("--top-k", type=int, default=0,help="top-k value to sample with")
    parser.add_argument("--temperature", type=float, default=1.02, help="temperature value to sample with")
    parser.add_argument("--top-p", type=float, default=1.0, help="top-p value to sample with")
    args = parser.parse_args()
    main(args)